package com.hank.websocketdemo.handler;

import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author gaoZhiHao
 * @date 2023/5/18 11:28
 */
@Component
public class ChatHandler extends TextWebSocketHandler {

    /**
     * 存session
     */
    private final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    /**
     * 存加入的用户
     */
    private final Map<String, String> userNames = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        sessions.put(sessionId, session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        JsonObject jsonObject = JsonParser.parseString(payload).getAsJsonObject();
        String type = jsonObject.get("type").getAsString();

        if ("join".equals(type)) {
            String userName = jsonObject.get("userName").getAsString();
            userNames.put(session.getId(), userName);

            JsonObject joinMessage = new JsonObject();
            joinMessage.addProperty("type", "join");
            joinMessage.addProperty("userName", userName);

            broadcast(joinMessage.toString());
        } else if ("chat".equals(type)) {
            String userName = userNames.get(session.getId());
            String text = jsonObject.get("text").getAsString();

            JsonObject chatMessage = new JsonObject();
            chatMessage.addProperty("type", "chat");
            chatMessage.addProperty("userName", userName);
            chatMessage.addProperty("text", text);

            broadcast(chatMessage.toString());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session.getId());

        String userName = userNames.get(session.getId());
        userNames.remove(session.getId());

        JsonObject leaveMessage = new JsonObject();
        leaveMessage.addProperty("type", "leave");
        leaveMessage.addProperty("userName", userName);

        broadcast(leaveMessage.toString());
    }

    private void broadcast(String message) {
        sessions.values().forEach(session -> {
            try {
                session.sendMessage(new TextMessage(message));
            } catch (IOException e) {
                e.printStackTrace();
            }
        });
    }
}
